package com.kingsoft.dc.khaos.dsl.spark.udf.custom

import com.kingsoft.dc.khaos.dsl.utils.IPUtil
import org.apache.spark.sql.catalyst.encoders.RowEncoder
import org.apache.spark.sql.types.StringType
import org.apache.spark.sql.{Column, DataFrame, Row}
import org.apache.spark.sql.functions._

import scala.collection.mutable.ArrayBuffer


/**
  * create by yansu on 2020/05/06 10:32
  */
class Ks3UdfCustom {

  /*
  0.udf分为两种实现方式 (1.spark udf 2.自己实现的逻辑)
  1.全局解析udf表达式;
  2.将表达式放到一个集合中,通过二分查找来判断是否存在于自定义方法中;
  3.
   */
  def get_IpState(data: DataFrame, ipFiled: String, targetFiled: String): DataFrame = {
    var res: DataFrame = data
    val flag = targetFiled + "_col_" + System.currentTimeMillis().toString
    val index = res.schema.fieldIndex(ipFiled)
    res = res.map(lines => {
      //      val rowWithSchema = new GenericRowWithSchema(lines.toSeq.toArray,res.schema)
      var remoteIpState = ""
      val ipUtil = IPUtil.getInstance()
      val ip = lines.get(index).toString
      if (null != ip) {
        val ipInfos = ipUtil.getRemoteIpInfos(ip)
        remoteIpState = ipInfos(0)
      }
      Row.merge(lines, Row(remoteIpState))
    })(RowEncoder(res.schema.add(flag, StringType)))
    filterColums(res, targetFiled, flag).withColumnRenamed(flag, targetFiled)
  }

  def get_IpProvince(data: DataFrame, ipFiled: String, targetFiled: String): DataFrame = {
    var res: DataFrame = data
    val flag = targetFiled + "_col_" + System.currentTimeMillis().toString
    val index = res.schema.fieldIndex(ipFiled)
    res = res.map(lines => {
      //      val rowWithSchema = new GenericRowWithSchema(lines.toSeq.toArray,res.schema)
      var remoteIpProvince = ""
      val ipUtil = IPUtil.getInstance()
      val ip = lines.get(index).toString
      if (null != ip) {
        val ipInfos = ipUtil.getRemoteIpInfos(ip)
        remoteIpProvince = ipInfos(1)
      }
      Row.merge(lines, Row(remoteIpProvince))
    })(RowEncoder(res.schema.add(flag, StringType)))
    filterColums(res, targetFiled, flag).withColumnRenamed(flag, targetFiled)
  }

  def get_IpCity(data: DataFrame, ipFiled: String, targetFiled: String): DataFrame = {
    var res: DataFrame = data
    val flag = targetFiled + "_col_" + System.currentTimeMillis().toString
    val index = res.schema.fieldIndex(ipFiled)
    res = res.map(lines => {
      //      val rowWithSchema = new GenericRowWithSchema(lines.toSeq.toArray,res.schema)
      var remoteIpCity = ""
      val ipUtil = IPUtil.getInstance()
      val ip = lines.get(index).toString
      if (null != ip) {
        val ipInfos = ipUtil.getRemoteIpInfos(ip)
        remoteIpCity = ipInfos(2)
      }
      Row.merge(lines, Row(remoteIpCity))
    })(RowEncoder(res.schema.add(flag, StringType)))
    res = filterColums(res, targetFiled, flag).withColumnRenamed(flag, targetFiled)
    res
  }

  def get_IpOperators(data: DataFrame, ipFiled: String, targetFiled: String): DataFrame = {
    var res: DataFrame = data
    val flag = targetFiled + "_col_" + System.currentTimeMillis().toString
    val index = res.schema.fieldIndex(ipFiled)
    res = res.map(lines => {
      //      val rowWithSchema = new GenericRowWithSchema(lines.toSeq.toArray,data.schema)
      var remoteIpIsp = "-"
      val ip = lines.get(index).toString
      val ipUtil = IPUtil.getInstance()
      if (null != ip) {
        val ipInfos = ipUtil.getRemoteIpInfos(ip)
        if (ipInfos(3).indexOf("/") == -1)
          remoteIpIsp = ipInfos(3)
        else
          remoteIpIsp = ipInfos(3).substring(0, ipInfos(3).indexOf("/"))
      }
      Row.merge(lines, Row(remoteIpIsp))
    })(RowEncoder(res.schema.add(flag, StringType)))
    filterColums(res, targetFiled, flag).withColumnRenamed(flag, targetFiled)
  }

  def filterColums(dataFrame: DataFrame, fieldName: String, flag: String): DataFrame = {
    val columnStrArray = dataFrame.columns
    val columnarArray = new ArrayBuffer[Column]()
    for (colStr <- columnStrArray) {
      if (!colStr.equalsIgnoreCase(fieldName)) {
        columnarArray += col(colStr)
      }
    }
    dataFrame.select(columnarArray: _*)
  }
}
